package com.zhidian.edm.util;

import com.zhidian.edm.entity.SysMenuPO;
import com.zhidian.edm.vo.RouterVO;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

public class TreeUtil {

    /**
     * 构建无根节点树，parentId 默认为 0
     *
     * @param list         所有节点
     * @param idGetter     获取id的方法
     * @param parentGetter 获取parentId的方法
     * @param <T>          节点类型（必须有 getChildren 方法）
     * @return 树形结构
     */
    public static <T> List<T> buildTree(
            List<T> list,
            Function<T, Long> idGetter,
            Function<T, Long> parentGetter
    ) {
        Map<Long, T> nodeMap = list.stream()
                .collect(Collectors.toMap(idGetter, Function.identity()));

        List<T> roots = new ArrayList<>();

        for (T node : list) {
            Long parentId = Optional.ofNullable(parentGetter.apply(node)).orElse(0L);
            if (parentId == 0) {
                roots.add(node);
            } else {
                T parent = nodeMap.get(parentId);
                if (parent != null) {
                    try {
                        // 通过反射获取 children 字段
                        var childrenField = parent.getClass().getDeclaredField("children");
                        childrenField.setAccessible(true);

                        List<T> children = (List<T>) childrenField.get(parent);
                        if (children == null) {
                            children = new ArrayList<>();
                            childrenField.set(parent, children);
                        }
                        children.add(node);
                    } catch (Exception e) {
                        throw new RuntimeException("构建树失败，请确保类中存在 List<T> children 字段", e);
                    }
                } else {
                    roots.add(node);
                }
            }
        }
        return roots;
    }

    public static List<RouterVO> buildRouterTree(List<SysMenuPO> menus, List<RouterVO> routers) {
        if (menus == null || routers == null || menus.size() != routers.size()) {
            throw new IllegalArgumentException("menus 和 routers 列表不能为空，且长度必须相同");
        }

        Map<Long, RouterVO> idRouterMap = new HashMap<>();
        for (int i = 0; i < routers.size(); i++) {
            idRouterMap.put(menus.get(i).getId(), routers.get(i));
        }

        List<RouterVO> tree = new ArrayList<>();
        for (int i = 0; i < routers.size(); i++) {
            SysMenuPO menu = menus.get(i);
            RouterVO router = routers.get(i);
            Long parentId = menu.getParentId() == null ? 0L : menu.getParentId();
            if (parentId == 0) {
                tree.add(router); // 根节点
            } else {
                RouterVO parent = idRouterMap.get(parentId);
                if (parent != null) {
                    if (parent.getChildren() == null) {
                        parent.setChildren(new ArrayList<>());
                    }
                    parent.getChildren().add(router);
                } else {
                    // 父节点不存在也算根
                    tree.add(router);
                }
            }
        }
        return tree;
    }
}
